# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：analysis.py
    Created on：2025/4/10 15:28
    Author：(input)
    Description:数据分析器
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class FlowAnalyzer:
    def __init__(self, df):
        """
        初始化分析器

        Parameters
        ----------
        df : pandas.DataFrame
            含有 x, y, phi, psi, magV, u, v, p, Cp 等列的流动数据
        exclusion_radius : float
            源汇位置的排除半径，单位与 x, y 相同
        """
        required_cols = {'x', 'y', 'phi', 'psi', 'magV', 'u', 'v', 'p', 'Cp'}
        if not required_cols.issubset(df.columns):
            raise ValueError(f"DataFrame 缺失必要列: {required_cols - set(df.columns)}")

        self.df = df.copy()

        # 推断网格尺寸
        self.nx = len(set(self.df['x']))
        self.ny = len(set(self.df['y']))
        self.X = self.df['x'].values.reshape((self.ny, self.nx))
        self.Y = self.df['y'].values.reshape((self.ny, self.nx))

        # 自动确定源和汇位置
        self._find_source_sink_positions()

        # 动态计算默认散点大小
        self.default_scatter_size = max(10, int(500 / (self.nx * self.ny)))  # 根据点密度调整大小

    def _find_source_sink_positions(self):
        """
        从 df 中找到源和汇的位置，结合 magV 和 psi 进行推断
        """
        # 步骤 1：找到 magV 的最大值和次大值位置
        sorted_by_magV = self.df.sort_values(by='magV', ascending=False)
        max_v_row = sorted_by_magV.iloc[0]  # 最大值
        second_max_v_row = sorted_by_magV.iloc[1]  # 次大值
        candidate_source = (max_v_row['x'], max_v_row['y'])
        candidate_sink = (second_max_v_row['x'], second_max_v_row['y'])

        # 步骤 2：结合 psi 验证源和汇
        # 源通常对应于 psi 的极大值，汇对应于 psi 的极小值
        sorted_by_psi = self.df.sort_values(by='psi', ascending=False)
        max_psi_row = sorted_by_psi.iloc[0]  # psi 最大值
        min_psi_row = sorted_by_psi.iloc[-1]  # psi 最小值
        psi_source = (max_psi_row['x'], max_psi_row['y'])
        psi_sink = (min_psi_row['x'], min_psi_row['y'])

        # 步骤 3：匹配 magV 和 psi 的位置
        # 计算候选位置与 psi 极值位置的距离
        dx = np.abs(self.X[0, 1] - self.X[0, 0])
        dy = np.abs(self.Y[1, 0] - self.Y[0, 0])
        tolerance = 2 * max(dx, dy)  # 容差为网格间距的 2 倍

        # 计算距离
        def distance(pos1, pos2):
            return np.sqrt((pos1[0] - pos2[0])**2 + (pos1[1] - pos2[1])**2)

        # 确定源位置
        if distance(candidate_source, psi_source) <= tolerance:
            self.source_pos = candidate_source
        elif distance(candidate_sink, psi_source) <= tolerance:
            self.source_pos = candidate_sink
        else:
            # 如果 magV 和 psi 位置不匹配，默认使用 magV 最大值
            self.source_pos = candidate_source

        # 确定汇位置
        if distance(candidate_sink, psi_sink) <= tolerance:
            self.sink_pos = candidate_sink
        elif distance(candidate_source, psi_sink) <= tolerance:
            self.sink_pos = candidate_source
        else:
            # 如果 magV 和 psi 位置不匹配，默认使用 magV 次大值
            self.sink_pos = candidate_sink

        # 调试信息
        print(f"自动确定的源位置: {self.source_pos}, magV = {self.df[(self.df['x'] == self.source_pos[0]) & (self.df['y'] == self.source_pos[1])]['magV'].iloc[0]:.3f}, psi = {self.df[(self.df['x'] == self.source_pos[0]) & (self.df['y'] == self.source_pos[1])]['psi'].iloc[0]:.3f}")
        print(f"自动确定的汇位置: {self.sink_pos}, magV = {self.df[(self.df['x'] == self.sink_pos[0]) & (self.df['y'] == self.sink_pos[1])]['magV'].iloc[0]:.3f}, psi = {self.df[(self.df['x'] == self.sink_pos[0]) & (self.df['y'] == self.sink_pos[1])]['psi'].iloc[0]:.3f}")

    def analyze_velocity_distribution(self, bins=5, plot=False, clip_range=None):
        """
        分析速度大小的分布范围，按分位数划分速度范围，并计算每个范围的点数占比

        Parameters
        ----------
        bins : int or list
            分位数数量或自定义的区间边界
        plot : bool
            是否绘制速度分布图
        clip_range : tuple or None
            速度大小裁剪的百分位范围，例如 (1, 99) 表示裁剪掉最低 1% 和最高 1% 的速度值，
            默认 None（不裁剪）

        Returns
        -------
        dict
            包含速度分布信息的字典
        """
        magV = self.df['magV'].values.copy()  # 使用原始数据

        # 双向裁剪速度大小
        v_lower = 0
        v_upper = np.inf
        if clip_range is not None:
            if not isinstance(clip_range, tuple) or len(clip_range) != 2:
                raise ValueError("clip_range 必须是一个包含两个元素的元组，例如 (1, 99)")
            lower_percentile, upper_percentile = clip_range
            if not (0 <= lower_percentile < upper_percentile <= 100):
                raise ValueError("clip_range 的值必须满足 0 <= lower_percentile < upper_percentile <= 100")
            v_lower = np.percentile(magV, lower_percentile)
            v_upper = np.percentile(magV, upper_percentile)
            magV = np.clip(magV, v_lower, v_upper)

        # 筛选出裁剪后仍有效的点
        mask = (magV >= v_lower) & (magV <= v_upper)
        magV_filtered = magV[mask]
        df_filtered = self.df[mask].copy()

        if len(magV_filtered) == 0:
            raise ValueError("裁剪后没有有效的速度数据，请调整 clip_range")

        if isinstance(bins, int):
            # 按分位数划分
            quantiles = np.linspace(0, 1, bins + 1)
            bin_edges = np.quantile(magV_filtered, quantiles)
            bin_edges = np.unique(bin_edges)  # 确保边界唯一
        else:
            bin_edges = bins

        # 计算每个区间的点数和占比
        hist, bin_edges = np.histogram(magV_filtered, bins=bin_edges)
        total_points = len(magV_filtered)
        proportions = hist / total_points

        # 构建分布信息
        distribution = []
        for i in range(len(hist)):
            range_str = f"[{bin_edges[i]:.3f}, {bin_edges[i + 1]:.3f})"
            points = hist[i]
            proportion = proportions[i]
            # 提取该速度范围内的点（基于裁剪后的速度）
            mask = (magV >= bin_edges[i]) & (magV < bin_edges[i + 1])
            subset = df_filtered[mask].copy()
            distribution.append({
                'range': range_str,
                'min_velocity': bin_edges[i],
                'max_velocity': bin_edges[i + 1],
                'points': points,
                'proportion': proportion,
                'subset': subset
            })

        if plot:
            # 绘制速度分布直方图
            plt.figure(figsize=(8, 6))
            plt.hist(magV_filtered, bins=bin_edges, density=True, alpha=0.7, color='skyblue', edgecolor='black')
            plt.xlabel('Velocity Magnitude (m/s)')
            plt.ylabel('Density')
            title = 'Velocity Magnitude Distribution'
            if clip_range is not None:
                title += f' (Clipped from {clip_range[0]}% to {clip_range[1]}%)'
            plt.title(title)
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.tight_layout()
            plt.show()

            # 绘制累积分布函数 (CDF)
            plt.figure(figsize=(8, 6))
            sorted_magV = np.sort(magV_filtered)
            cdf = np.arange(1, len(sorted_magV) + 1) / len(sorted_magV)
            plt.plot(sorted_magV, cdf, color='navy')
            plt.xlabel('Velocity Magnitude (m/s)')
            plt.ylabel('Cumulative Probability')
            title = 'Cumulative Distribution of Velocity Magnitude'
            if clip_range is not None:
                title += f' (Clipped from {clip_range[0]}% to {clip_range[1]}%)'
            plt.title(title)
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.tight_layout()
            plt.show()

        return {'bin_edges': bin_edges, 'distribution': distribution}

    def plot_velocity_distribution_regions(self, bins=5, clip_range=(1, 99), scatter_size=None):
        """
        绘制不同速度大小范围的区域分布，裁剪速度以避免奇点影响和速度为零的部分

        Parameters
        ----------
        bins : int or list
            分位数数量或自定义的区间边界
        clip_range : tuple
            速度大小裁剪的百分位范围，例如 (1, 99) 表示裁剪掉最低 1% 和最高 1% 的速度值，
            默认 (1, 99)
        scatter_size : float or None
            散点图中圆点的大小，如果为 None，则使用动态计算的默认值
        """
        dist_info = self.analyze_velocity_distribution(
            bins=bins, plot=False, clip_range=clip_range
        )
        bin_edges = dist_info['bin_edges']
        distribution = dist_info['distribution']

        # 如果未指定 scatter_size，使用动态计算的默认值
        if scatter_size is None:
            scatter_size = self.default_scatter_size

        plt.figure(figsize=(8, 6))
        for i, region in enumerate(distribution):
            subset = region['subset']
            # 在图例中显示速度范围和占比
            label = f"Velocity {region['range']} ({region['proportion'] * 100:.1f}%)"
            plt.scatter(subset['x'], subset['y'], s=scatter_size, label=label, alpha=0.6)

        # 标注源和汇位置
        if self.source_pos is not None:
            plt.scatter(self.source_pos[0], self.source_pos[1], c='red', marker='o', s=100, label='Source')
        if self.sink_pos is not None:
            plt.scatter(self.sink_pos[0], self.sink_pos[1], c='green', marker='x', s=100, label='Sink')

        plt.xlabel('x')
        plt.ylabel('y')
        title = 'Regions of Different Velocity Ranges'
        if clip_range is not None:
            title += f' (Clipped from {clip_range[0]}% to {clip_range[1]}%)'
        plt.title(title)
        plt.legend()
        plt.axis('equal')
        plt.tight_layout()
        plt.show()